import numpy as np

from fund.strategy.strategy_util import process_dates, empty_result
from util.csv_util import csv_2_df
from util.date_util import *

def strategy_based_on_drop(df, money=200, start_date=None, end_date=None):
    """
    多日跌幅触发的买入策略（加强版）
    核心逻辑：根据连续下跌天数对应的跌幅阈值来触发买入
    跌幅阈值数组：[1天-3%, 2天-5%, 3天-7%, 4天-8.5%, 5天-10%]
    
    :param df: 基金净值数据（需包含日期索引和'单位净值'列）
    :param money: 每次买入金额
    :param start_date: 策略开始日期
    :param end_date: 策略结束日期
    :return: 包含收益指标和买入记录的字典
    """
    # # --- 固定跌幅阈值数组 ---
    drop_thresholds = [
        (1, 0.030),
        (2, 0.050),
        (3, 0.070),
        (4, 0.085),
        (5, 0.100)
    ]

    # drop_thresholds = [
    #     (1, 0.020),
    #     (2, 0.038),
    #     (3, 0.054),
    #     (4, 0.068),
    #     (5, 0.080)
    # ]

    # drop_thresholds = [
    #     (1, 0.030),
    #     (2, 0.055),
    #     (3, 0.075),
    #     (4, 0.090),
    #     (5, 0.100)
    # ]

    # --- 初始化 ---
    start_date, end_date = process_dates(df, start_date, end_date)
    total_cost = 0
    total_units = 0
    buy_times = 0

    # 记录买入操作
    buy_records = []

    # 获取日期索引位置
    start_index = df.index.searchsorted(start_date)
    end_index = df.index.searchsorted(end_date)

    # 确保有足够数据计算最大天数跌幅（需要前5天的数据）
    start_index = max(start_index, 5)

    if start_index >= end_index:
        return empty_result(start_date, end_date)

    # --- 核心逻辑 ---
    for current_pos in range(start_index, end_index + 1):
        current_date = df.index[current_pos]
        current_net_value = df['单位净值'].iloc[current_pos]

        trigger_reason = ""
        max_days_met = 0
        max_drop_ratio = 0

        # 检查所有跌幅条件
        for days, threshold in drop_thresholds:
            # 检查是否有足够的历史数据
            if current_pos < days:
                continue

            # 计算N日跌幅
            prev_nday_value = df['单位净值'].iloc[current_pos - days]
            drop_nday_ratio = (prev_nday_value - current_net_value) / prev_nday_value

            # 记录满足条件的最大跌幅和天数
            if drop_nday_ratio >= threshold:
                if days > max_days_met:
                    max_days_met = days
                    max_drop_ratio = drop_nday_ratio
                    trigger_reason = f"{days}日跌幅{drop_nday_ratio:.2%}≥{threshold:.1%}"

        # 如果满足任一条件就买入（选择最大天数的条件）
        if max_days_met > 0:
            daily_units = money / current_net_value
            total_units += daily_units
            total_cost += money
            buy_times += 1

            # 记录买入操作详情
            buy_records.append({
                'date': current_date,
                'price': current_net_value,
                'cost': money,
                'units': daily_units,
                'trigger_reason': trigger_reason,
                'drop_days': max_days_met,
                'drop_ratio': max_drop_ratio
            })

    # --- 收益计算 ---
    if end_index > start_index:
        last_net_value = df['单位净值'].iloc[end_index - 1]
        final_asset = total_units * last_net_value
    else:
        final_asset = 0

    total_profit = final_asset - total_cost
    total_return_rate = total_profit / total_cost if total_cost != 0 else 0
    investment_days = (end_date - start_date).days

    # 年化收益率
    if total_cost != 0 and investment_days > 0:
        annual_return_rate = (final_asset / total_cost) ** (365 / investment_days) - 1
    else:
        annual_return_rate = 0

    return {
        "开始时间": start_date.strftime('%Y-%m-%d'),
        "结束时间": end_date.strftime('%Y-%m-%d'),
        "总投入": total_cost,
        "最终资产": final_asset,
        "总利润": total_profit,
        "总收益率": total_return_rate,
        "年化收益率": annual_return_rate,
        "买入次数": buy_times,
        "买入记录": buy_records,
    }


def calculate_daily_return_distribution(
        df,
        start_date=None,
        end_date=None,
        return_ranges=None
):
    """
    统计指定时间范围内基金每日涨跌幅的区间分布比例

    :param df: 基金净值数据（需包含日期索引和'单位净值'列）
    :param start_date: 统计开始日期（字符串，如"20220101"）
    :param end_date: 统计结束日期（字符串，如"20250101"）
    :param return_ranges: 自定义涨跌幅区间，默认提供常见区间
    :return: 字典，key=区间名称，value={count:天数, ratio:占比(%)}
    """
    # ---------------------- 1. 初始化默认涨跌幅区间（左闭右开） ----------------------
    if return_ranges is None:
        return_ranges = [
            # 下跌区间：从“跌5%及以上”到“跌0%-1%”，每1%1个区间
            ("大跌（≤-5%）", -np.inf, -0.05),  # 跌幅≥5%（含5%）
            ("跌4%-5%", -0.05, -0.04),  # 跌幅4%~5%（含4%，不含5%）
            ("跌3%-4%", -0.04, -0.03),  # 跌幅3%~4%
            ("跌2%-3%", -0.03, -0.02),  # 跌幅2%~3%
            ("跌1%-2%", -0.02, -0.01),  # 跌幅1%~2%
            ("跌0%-1%", -0.01, 0),  # 跌幅0%~1%（接近持平的小跌）

            # 持平区间：微小波动（±0.1%内）单独归类
            ("持平（±0.1%内）", -0.001, 0.001),

            # 上涨区间：从“涨0%-1%”到“涨5%及以上”，每1%1个区间
            ("涨0%-1%", 0.001, 0.01),  # 涨幅0%~1%（接近持平的小涨）
            ("涨1%-2%", 0.01, 0.02),  # 涨幅1%~2%（你需求中的核心区间）
            ("涨2%-3%", 0.02, 0.03),  # 涨幅2%~3%
            ("涨3%-4%", 0.03, 0.04),  # 涨幅3%~4%
            ("涨4%-5%", 0.04, 0.05),  # 涨幅4%~5%
            ("大涨（≥5%）", 0.05, np.inf)  # 涨幅≥5%（含5%）
        ]

    # ---------------------- 2. 筛选指定时间范围的净值数据 ----------------------
    # 处理开始/结束日期（复用你已有的process_dates，确保日期格式统一）
    start_date_dt, end_date_dt = process_dates(df, start_date, end_date)
    # 截取目标时间范围内的数据
    df_filtered = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)].copy()

    # ---------------------- 3. 计算每日涨跌幅（关键步骤） ----------------------
    # 涨跌幅 = (当日单位净值 - 前一日单位净值) / 前一日单位净值
    df_filtered["daily_return"] = df_filtered["单位净值"].pct_change()
    # 过滤掉第一天（pct_change后第一天为NaN，无涨跌幅数据）
    valid_returns = df_filtered["daily_return"].dropna()

    # 处理无有效数据的情况
    if len(valid_returns) == 0:
        return {"error": "指定时间范围内无有效涨跌幅数据（可能数据不足1天）"}

    # ---------------------- 4. 统计每个区间的天数和占比 ----------------------
    distribution_result = {}
    total_valid_days = len(valid_returns)  # 有效交易日总数

    # 遍历每个涨跌幅区间，统计数量和占比
    for range_name, left_bound, right_bound in return_ranges:
        # 统计该区间内的涨跌幅天数（左闭右开：>=left_bound 且 <right_bound）
        count = len(valid_returns[(valid_returns >= left_bound) & (valid_returns < right_bound)])
        # 计算占比（保留2位小数）
        ratio = (count / total_valid_days) * 100
        # 存入结果字典
        distribution_result[range_name] = {
            "count": count,  # 该区间的天数
            "ratio": round(ratio, 2)  # 该区间的占比（%）
        }

    # 添加统计汇总信息
    distribution_result["汇总"] = {
        "统计时间范围": f"{start_date_dt.strftime('%Y-%m-%d')} ~ {end_date_dt.strftime('%Y-%m-%d')}",
        "有效交易日总数": total_valid_days,
        "所有区间占比之和": round(sum([v["ratio"] for v in distribution_result.values() if "ratio" in v]), 2)
    }

    return distribution_result


# ---------------------- 示例：如何使用这个函数 ----------------------
if __name__ == "__main__":
    # 1. 读取基金数据
    # fund_code = '022364'
    # fund_code = '580006'
    fund_code = '018124'
    df = csv_2_df(f"fund/history_data/{fund_code}-daily.csv")

    res = strategy_based_on_drop(df, start_date=before_2y_yyyymmdd())
    print(res['买入记录'])

    # # 2. 统计近3年的涨跌幅分布（和你之前的策略时间范围一致）
    # start_date = "20220101"  # 或用before_n_years_yyyymmdd(3)
    # end_date = "20230101"  # 或用yesterday_yyyymmdd()
    # result = calculate_daily_return_distribution(df, start_date=start_date, end_date=end_date)

    # # 3. 打印结果（格式化输出，更易读）
    # print(f"涨跌幅分布统计 [{fund_code}]：")
    # print(f"统计区间：{result['汇总']['统计时间范围']}")
    # print(f"有效交易日总数：{result['汇总']['有效交易日总数']}天\n")

    # # 打印每个区间的详情（排除“汇总”项）
    # for range_name, data in result.items():
    #     if range_name != "汇总" and "error" not in result:
    #         print(f"{range_name:15s}：{data['count']:3d}天（占比{data['ratio']:5.2f}%）")

    # # 打印错误信息（如有）
    # if "error" in result:
    #     print(f"\n错误：{result['error']}")
